# 采用阿里云百练Api模式
from ast import increment_lineno

from dashscope import Generation
import dashscope
from agent.qwen.api.ApiModel import ApiModel

class DashScopeModel(ApiModel):
    def __init__(self, api_key, base_url, model, stream, thinking):
        super().__init__(api_key, base_url, model, stream, thinking)
        dashscope.api_key = api_key
        dashscope.base_url = base_url


    # 输出Token数量
    def print_token_count(self, token_usage):
        if token_usage is not None:
            print(f"\n--- 请求用量, 输入 Tokens: {token_usage.input_tokens}， 输出 Tokens: {token_usage.output_tokens}, 总计 Tokens: {token_usage.total_tokens}")

    # 非流式输出
    def print_text_not_stream(self, resp) -> str:
        answer = resp.output.choices[0].message.content
        print(f"模型：{answer}")
        return answer

    # 流式输出
    def print_text_stream(self, resp) -> str:
        answer = ""
        is_answering = False
        # 深度思考
        if self.thinking:
            print("=" * 20 + "思考过程" + "=" * 20)
        for chunk in resp:
            # print(chunk)
            # 深度思考内容
            if chunk.output.choices[0].message.reasoning_content != "":
                 if not is_answering:
                     print(chunk.output.choices[0].message.reasoning_content, end="", flush=True)
            # 回复内容
            elif  chunk.output.choices[0].message.content != "":
                if not is_answering:
                    print("\n" + "=" * 20 + "完整回复" + "=" * 20)
                    is_answering = True

                content = chunk.output.choices[0].message.content
                print(content, end="", flush=True)
                answer += content
            # 输出Token数量
            elif chunk.output.choices[0].finish_reason == "stop":
                self.print_token_count(chunk.usage)
        return answer


    def chat(self, messages) ->str:
        try:
            resp = Generation.call(self.model,
                                   messages=messages,
                                   # 是否流式输出
                                   stream=self.stream,
                                   # 是否深度思考
                                   enable_thinking=self.thinking,
                                   # message格式
                                   result_format="message",
                                   # 增量输出
                                   incremental_output=True)

        except Exception as e:
            print(e)
            return "系统繁忙，请稍后再试"
        else:
            return super().print_text_answer(resp)
    